/*
 * mahony.c
 *
 *  Created: 20250116
 *  Author: Hecate
 */
#include "mahony.h"

#include "lowPassFilter.h"

#include <math.h>

//---------------------------------------------------------------------------------------------------
// Definitions

#define sampleFreq 1000.0f      // sample frequency in Hz
#define twoKpDef (2.0f * 0.5f)  // 2 * proportional gain
#define twoKiDef (2.0f * 0.0f)  // 2 * integral gain

//---------------------------------------------------------------------------------------------------
// Variable definitions

static volatile float twoKp = twoKpDef;                                            // 2 * proportional gain (Kp)
static volatile float twoKi = twoKiDef;                                            // 2 * integral gain (Ki)
static volatile float integralFBx = 0.0f, integralFBy = 0.0f, integralFBz = 0.0f;  // integral error terms scaled by Ki

//---------------------------------------------------------------------------------------------------
// Function declarations

static inline float invSqrt(float x);

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

static void MahonyAHRSupdate(float q[4], float gx, float gy, float gz, float ax, float ay, float az, float mx, float my, float mz);
static void MahonyAHRSupdateIMU(float q[4], float gx, float gy, float gz, float ax, float ay, float az);

static LowPassFilterData* yawFilter;

void mahony6Init(_6IMUDATA* obj, float* gx, float* gy, float* gz, float* ax, float* ay, float* az)
{
    obj->q[0] = 1.0f;
    obj->q[1] = 0.0f;
    obj->q[2] = 0.0f;
    obj->q[3] = 0.0f;
    obj->gx = gx;
    obj->gy = gy;
    obj->gz = gz;
    obj->ax = ax;
    obj->ay = ay;
    obj->az = az;

    obj->pitch = 0.0f;
    obj->roll = 0.0f;
    obj->yaw = 0.0f;

    yawFilter = lowPassFilterInit(0.1f);
}
void mahony9Init(_9IMUDATA* obj, float* gx, float* gy, float* gz, float* ax, float* ay, float* az, float* mx, float* my, float* mz)
{
    obj->q[0] = 1.0f;
    obj->q[1] = 0.0f;
    obj->q[2] = 0.0f;
    obj->q[3] = 0.0f;
    obj->gx = gx;
    obj->gy = gy;
    obj->gz = gz;
    obj->ax = ax;
    obj->ay = ay;
    obj->az = az;
    obj->mx = mx;
    obj->my = my;
    obj->mz = mz;

    obj->pitch = 0.0f;
    obj->roll = 0.0f;
    obj->yaw = 0.0f;
}

static inline void quaternionToEuler(const float q[4], float* pitch, float* roll, float* yaw);
static volatile const float yawfra = (1.0f / sampleFreq);

void mahony6Update(_6IMUDATA* obj)
{
    MahonyAHRSupdateIMU(obj->q, *(obj->gx), *(obj->gy), *(obj->gz), *(obj->ax), *(obj->ay), *(obj->az));

    lowPassFilter(yawFilter, *(obj->gz));

    quaternionToEuler(obj->q, &(obj->pitch), &(obj->roll), NULL);
    obj->yaw += lowPassFilterGetValue(yawFilter) * yawfra;
}
void mahony9Update(_9IMUDATA* obj)
{
    MahonyAHRSupdate(obj->q, *(obj->gx), *(obj->gy), *(obj->gz), *(obj->ax), *(obj->ay), *(obj->az), *(obj->mx), *(obj->my), *(obj->mz));

    quaternionToEuler(obj->q, &(obj->pitch), &(obj->roll), &(obj->yaw));
}

//====================================================================================================
// Functions

//---------------------------------------------------------------------------------------------------
// AHRS algorithm update

static void MahonyAHRSupdate(float q[4], float gx, float gy, float gz, float ax, float ay, float az, float mx, float my, float mz)
{
    float recipNorm;
    float q0q0, q0q1, q0q2, q0q3, q1q1, q1q2, q1q3, q2q2, q2q3, q3q3;
    float hx, hy, bx, bz;
    float halfvx, halfvy, halfvz, halfwx, halfwy, halfwz;
    float halfex, halfey, halfez;
    float qa, qb, qc;

    // Use IMU algorithm if magnetometer measurement invalid (avoids NaN in magnetometer normalisation)
    if ((mx == 0.0f) && (my == 0.0f) && (mz == 0.0f))
    {
        MahonyAHRSupdateIMU(q, gx, gy, gz, ax, ay, az);
        return;
    }

    // Compute feedback only if accelerometer measurement valid (avoids NaN in accelerometer normalisation)
    if (!((ax == 0.0f) && (ay == 0.0f) && (az == 0.0f)))
    {
        // Normalise accelerometer measurement
        recipNorm = invSqrt(ax * ax + ay * ay + az * az);
        ax *= recipNorm;
        ay *= recipNorm;
        az *= recipNorm;

        // Normalise magnetometer measurement
        recipNorm = invSqrt(mx * mx + my * my + mz * mz);
        mx *= recipNorm;
        my *= recipNorm;
        mz *= recipNorm;

        // Auxiliary variables to avoid repeated arithmetic
        q0q0 = q[0] * q[0];
        q0q1 = q[0] * q[1];
        q0q2 = q[0] * q[2];
        q0q3 = q[0] * q[3];
        q1q1 = q[1] * q[1];
        q1q2 = q[1] * q[2];
        q1q3 = q[1] * q[3];
        q2q2 = q[2] * q[2];
        q2q3 = q[2] * q[3];
        q3q3 = q[3] * q[3];

        // Reference direction of Earth's magnetic field
        hx = 2.0f * (mx * (0.5f - q2q2 - q3q3) + my * (q1q2 - q0q3) + mz * (q1q3 + q0q2));
        hy = 2.0f * (mx * (q1q2 + q0q3) + my * (0.5f - q1q1 - q3q3) + mz * (q2q3 - q0q1));
        bx = sqrtf(hx * hx + hy * hy);
        bz = 2.0f * (mx * (q1q3 - q0q2) + my * (q2q3 + q0q1) + mz * (0.5f - q1q1 - q2q2));

        // Estimated direction of gravity and magnetic field
        halfvx = q1q3 - q0q2;
        halfvy = q0q1 + q2q3;
        halfvz = q0q0 - 0.5f + q3q3;
        halfwx = bx * (0.5f - q2q2 - q3q3) + bz * (q1q3 - q0q2);
        halfwy = bx * (q1q2 - q0q3) + bz * (q0q1 + q2q3);
        halfwz = bx * (q0q2 + q1q3) + bz * (0.5f - q1q1 - q2q2);

        // Error is sum of cross product between estimated direction and measured direction of field vectors
        halfex = (ay * halfvz - az * halfvy) + (my * halfwz - mz * halfwy);
        halfey = (az * halfvx - ax * halfvz) + (mz * halfwx - mx * halfwz);
        halfez = (ax * halfvy - ay * halfvx) + (mx * halfwy - my * halfwx);

        // Compute and apply integral feedback if enabled
        if (twoKi > 0.0f)
        {
            integralFBx += twoKi * halfex * (1.0f / sampleFreq);  // integral error scaled by Ki
            integralFBy += twoKi * halfey * (1.0f / sampleFreq);
            integralFBz += twoKi * halfez * (1.0f / sampleFreq);
            gx += integralFBx;  // apply integral feedback
            gy += integralFBy;
            gz += integralFBz;
        }
        else
        {
            integralFBx = 0.0f;  // prevent integral windup
            integralFBy = 0.0f;
            integralFBz = 0.0f;
        }

        // Apply proportional feedback
        gx += twoKp * halfex;
        gy += twoKp * halfey;
        gz += twoKp * halfez;
    }

    // Integrate rate of change of quaternion
    gx *= (0.5f * (1.0f / sampleFreq));  // pre-multiply common factors
    gy *= (0.5f * (1.0f / sampleFreq));
    gz *= (0.5f * (1.0f / sampleFreq));
    qa = q[0];
    qb = q[1];
    qc = q[2];
    q[0] += (-qb * gx - qc * gy - q[3] * gz);
    q[1] += (qa * gx + qc * gz - q[3] * gy);
    q[2] += (qa * gy - qb * gz + q[3] * gx);
    q[3] += (qa * gz + qb * gy - qc * gx);

    // Normalise quaternion
    recipNorm = invSqrt(q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]);
    q[0] *= recipNorm;
    q[1] *= recipNorm;
    q[2] *= recipNorm;
    q[3] *= recipNorm;
}

//---------------------------------------------------------------------------------------------------
// IMU algorithm update

static void MahonyAHRSupdateIMU(float q[4], float gx, float gy, float gz, float ax, float ay, float az)
{
    float recipNorm;
    float halfvx, halfvy, halfvz;
    float halfex, halfey, halfez;
    float qa, qb, qc;

    static float half_errorInv;
    static int count = 0;
    /********************************参数规划*********************************************************************************************************************/

    if (count < 3000)
    {
        count += 2;
        twoKp = 20;
    }
    else
    {
        if (half_errorInv < 20)
            twoKp = 0.1f;
        else
            twoKp = 4;
    }

    /*****************************************************************************************************************************************************/

    // Compute feedback only if accelerometer measurement valid (avoids NaN in accelerometer normalisation)
    if (!((ax == 0.0f) && (ay == 0.0f) && (az == 0.0f)))
    {
        // Normalise accelerometer measurement
        recipNorm = invSqrt(ax * ax + ay * ay + az * az);
        ax *= recipNorm;
        ay *= recipNorm;
        az *= recipNorm;

        // Estimated direction of gravity and vector perpendicular to magnetic flux
        halfvx = q[1] * q[3] - q[0] * q[2];
        halfvy = q[0] * q[1] + q[2] * q[3];
        halfvz = q[0] * q[0] - 0.5f + q[3] * q[3];

        // Error is sum of cross product between estimated and measured direction of gravity
        halfex = (ay * halfvz - az * halfvy);
        halfey = (az * halfvx - ax * halfvz);
        halfez = (ax * halfvy - ay * halfvx);
        half_errorInv = invSqrt(halfex * halfex + halfey * halfey + halfez * halfez);
        // Compute and apply integral feedback if enabled
        if (twoKi > 0.0f)
        {
            integralFBx += twoKi * halfex * (1.0f / sampleFreq);  // integral error scaled by Ki
            integralFBy += twoKi * halfey * (1.0f / sampleFreq);
            integralFBz += twoKi * halfez * (1.0f / sampleFreq);
            gx += integralFBx;  // apply integral feedback
            gy += integralFBy;
            gz += integralFBz;
        }
        else
        {
            integralFBx = 0.0f;  // prevent integral windup
            integralFBy = 0.0f;
            integralFBz = 0.0f;
        }

        // Apply proportional feedback
        gx += twoKp * halfex;
        gy += twoKp * halfey;
        gz += twoKp * halfez;
    }

    // Integrate rate of change of quaternion
    gx *= (0.5f * (1.0f / sampleFreq));  // pre-multiply common factors
    gy *= (0.5f * (1.0f / sampleFreq));
    gz *= (0.5f * (1.0f / sampleFreq));
    qa = q[0];
    qb = q[1];
    qc = q[2];
    q[0] += (-qb * gx - qc * gy - q[3] * gz);
    q[1] += (qa * gx + qc * gz - q[3] * gy);
    q[2] += (qa * gy - qb * gz + q[3] * gx);
    q[3] += (qa * gz + qb * gy - qc * gx);

    // Normalise quaternion
    recipNorm = invSqrt(q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]);
    q[0] *= recipNorm;
    q[1] *= recipNorm;
    q[2] *= recipNorm;
    q[3] *= recipNorm;
}

//---------------------------------------------------------------------------------------------------
// Fast inverse square-root
// See: http://en.wikipedia.org/wiki/Fast_inverse_square_root

static inline float invSqrt(float x)
{
    float halfx = 0.5f * x;
    float y = x;
    long i = *(long*) &y;
    i = 0x5f3759df - (i >> 1);
    y = *(float*) &i;
    y = y * (1.5f - (halfx * y * y));
    return y;
}

static inline void quaternionToEuler(const float q[4], float* pitch, float* roll, float* yaw)
{
    if (q == NULL)
        return;

    const float q0 = q[0];  // w
    const float q1 = q[1];  // x
    const float q2 = q[2];  // y
    const float q3 = q[3];  // z

    const float sinr_cosp = 2.0f * (q0 * q1 + q2 * q3);
    const float cosr_cosp = 1.0f - 2.0f * (q1 * q1 + q2 * q2);
    *roll = atan2f(sinr_cosp, cosr_cosp);

    float sinp = 2.0f * (q0 * q2 - q1 * q3);
    if (sinp > 1.0f)
        sinp = 1.0f;
    else if (sinp < -1.0f)
        sinp = -1.0f;
    *pitch = asinf(sinp);

    if (yaw == NULL) return;
    const float siny_cosp = 2.0f * (q0 * q3 + q1 * q2);
    const float cosy_cosp = 1.0f - 2.0f * (q2 * q2 + q3 * q3);
    *yaw = atan2f(siny_cosp, cosy_cosp);
}
